from vBaum import *
from math import *

from tree_common import *

# Monopodial: Main axis dominates
class MonopodialBranch:
	def __init__(self):
		self.branchAngle = radians(68)
		self.branchAngleSigma = radians(10)
		self.numBranchesPerNode = 4.5
		self.numBranchesPerNodeRange = 0.4 #sigma * numBranchesPerNode
		
		self.angleBetweenBranchesAtNodeSigma = radians(10)
		self.divergence = True
		self.successor = self
		self.orientUp = False
		self.radiusFallOff = 0.88 # 0.5

		self.roundNodes = True
		self.nodeRandomness = 0.05 #per one unit branch length
		self.branchRoughness = 0.6
		self.initialInterNodeLength = 0.15
		self.interNodeLengthGrowth = 0.04
		self.interNodePoints = 3 # sets number of spline points on internodes
		self.angleOffsetSigma = 0.1
		self.minBranchDiameter = voxelSize()
		self.minBranchLength = max(self.initialInterNodeLength*0.75, voxelSize()*7)
		
		self.branchRadiusScale = 1.5
		
		#used to calculate RadiusHeightConstant
		designLength = 20
		designRadius = 0.25
		
		self.RadiusHeightConstant = calcRadiusLengthConstant(designLength, designRadius, self.radiusFallOff)
		
	def heightToDiam(self, length):
		return LengthToDiam(length, self.RadiusHeightConstant, self.radiusFallOff)
				
	def generateGeometry(self, startPos, direction, trunkRadius):
		#for graphics output, don't generate small diameters
		graphicsTrunkRadius = max(self.minBranchDiameter, trunkRadius)
		
		#penalize small branches a bit
		maxLength = diamToLength(trunkRadius, self.RadiusHeightConstant, self.radiusFallOff)
		
		#reject too small branches
		if(maxLength < self.minBranchLength):
			return Volume()
		
		print("Max length is ", maxLength)
				
		endPos = startPos + direction * maxLength
		
		nodeHeights = generateNodePositions(maxLength, self.initialInterNodeLength, self.interNodeLengthGrowth)
		nodePoints, branchContour = generateNodesOnBranch(startPos, endPos, nodeHeights, self.interNodePoints, self.branchRoughness, self.nodeRandomness)

		print("length of nodeHeights is", len(nodeHeights))
		

		actualLength = nodeHeights[-1]
		print("actual length is ", actualLength)
		
		print("generated", len(nodePoints), "nodes.")
		print("trunk radius:", trunkRadius)
		
		#create this branch
#		print("branchPositions:")
#		for position in nodePoints:
#			printVec(position)
			
		print("branchContour:")
		for position in branchContour:
			printVec(position)
			
#		branch = createBranch(branchContour, graphicsTrunkRadius * self.branchRadiusScale, voxelSize())
		branch = createBranch(branchContour, graphicsTrunkRadius * self.branchRadiusScale, voxelSize())
		
		print("finshed creating branch")
		
		angleOffset = 0
		for i in range(1, len(nodePoints)-1):
			#calculate area loss of this internode
			currRadius = (1 - nodeHeights[i] / actualLength) * trunkRadius
			lastRadius = (1 - nodeHeights[i-1] / actualLength) * trunkRadius
			areaLoss = lastRadius*lastRadius - currRadius*currRadius
			
			numBranches = int(rand.gauss(0, self.numBranchesPerNode * self.numBranchesPerNodeRange)) + self.numBranchesPerNode
			
			numBranches = max(int(numBranches), 0)
			
			if numBranches <= 0:
				continue;
			
			branchArea = areaLoss / numBranches
			nodeRadius = sqrt(branchArea)
			
			nodePos = nodePoints[i]
			
			print("nodeRadius =", nodeRadius)
			angleStep = (2 * pi) / (numBranches)
			angle = 0
			if self.divergence: angle = angleOffset
	
			for a in range(0, numBranches):
				firstBranchDir = computeFirstBranchDir(direction, rand.gauss(self.branchAngle, self.branchAngleSigma))
				if self.orientUp:
					firstBranchDir = rotateDir(direction, Vec3(0, 1, 0), rand.gauss(self.branchAngle, self.branchAngleSigma))
				angle = angle + angleStep
				if angle > 2 * pi: angle = angle - (2 * pi)
				branchDir = rotateDir(firstBranchDir, direction, rand.gauss(angle, self.angleBetweenBranchesAtNodeSigma))
				subBranch = self.successor.generateGeometry(nodePos, branchDir, nodeRadius)
				branch = branch.union(subBranch)

			#DEBUG
#			sph = Sphere(lastRadius * 5 * voxelsPerUnit)
#			sph.setPosition(nodePos * voxelsPerUnit)
#			branch = branch.union(sph)
			
			angleOffset = angleOffset + angleStep * (0.5 + rand.gauss(0, self.angleOffsetSigma))
			if angleOffset > 2 * pi: angleOffset = angleOffset - (2 * pi)
			
			if self.roundNodes: branch = roundSubVolume(branch, nodePos, nodeRadius*2.0, nodeRadius*0.5)
			
		return branch


node = MonopodialBranch()
node2 = MonopodialBranch()
node2.numBranchesPerNode = 2
node2.numBranchesPerNodeRange = 0
node2.nodeRandomness = 0.08
#node2.RadiusHeightConstant
node2.orientUp = True
node2.orientNodes = False
node2.divergence = False
node.successor = node2
tree = node.generateGeometry(Vec3(0, 0, 0), Vec3(0, 1, 0), node.heightToDiam(5))
tree.exportAsObj("monopodialTree", voxelsPerUnit, 2)


